package dan.utility
import chisel3._
import chisel3.util._
import dan.common._
import scala.language.implicitConversions
import dan.common.Consts.UOP_X
import dan.frontend.BrUpdateInfo
import dan.common.Consts.IMM_U5
import dan.common.Consts.IMM_U12
import dan.common.Consts.IMM_S12
import dan.common.Consts.IMM_S14
import dan.common.Consts.IMM_S16
import dan.common.Consts.IMM_U20
import dan.common.Consts.IMM_S20
import dan.common.Consts.IMM_S26
import dan.common.Consts.IMM_CSR

object IsEqual{
    def apply(a: UInt, b: UInt): Bool = {
        !(a ^ b).orR
    }
}

// 从input的第一个1开始，低位方向的全1掩码
object MaskLower {
    def apply(input: UInt): UInt = {
        val n = input.getWidth
        (0 until n).map(i => input >> i.U).reduce(_|_)
    }
}

// 从input的第一个1开始，高位方向的全1掩码
object MaskUpper {
    def apply(in: UInt):UInt = {
        val n = in.getWidth
        (0 until n).map(i => (in << i.U)(n-1,0)).reduce(_|_)
    }
}

// 循环计数，递增
object WrapInc {
    def apply(curVal: UInt, upLimit: Int): UInt = {
        // 
        if (isPow2(upLimit)) {
            (curVal + 1.U)(log2Ceil(upLimit) - 1, 0)
        } else {
            val wrap = curVal === (upLimit - 1).U
            Mux(wrap, 0.U, curVal + 1.U)
        }
    }
}

// 递减
object WrapDec {
    def apply(curVal: UInt, upLimit: Int): UInt = {
        if (isPow2(upLimit)) {
            (curVal - 1.U)(log2Ceil(upLimit) - 1,0)
        } else {
            val wrap = curVal === 0.U
            Mux(wrap, (upLimit-1).U, curVal - 1.U)
        }
    }
}

// 加法
object WrapAdd {
    def apply(curVal: UInt, inc: UInt, upLimit: Int): UInt = {
        if (isPow2(upLimit)) {
            (curVal + inc)(log2Ceil(upLimit) - 1, 0)
        } else {
            // 防止溢出
            val padinVal = curVal.pad(curVal.getWidth + 1)
            val result = padinVal + inc
            val wrap = result >= upLimit.U
            Mux(wrap, result - upLimit.U, result)
        }
    }
}

// 判断环形队列中idx1和idx2的元素年龄
/* 
    0 - (head-1)为“新区”，head之后为“旧区”
    1 ^ 1 ^ 1，都属于新区
    1 ^ 0 ^ 0，都属于旧区
    0 ^ 0 ^ 1，idx1属于旧区，idx2属于新区
 */
object IsOlder{
    def apply(idx1: UInt, idx2: UInt, head: UInt): Bool = (idx1 < idx2) ^ (idx1 < head) ^ (idx2 < head)
}


// 存在一位重叠1
object maskMatch{
    def apply(mask1: UInt, mask2: UInt): Bool = (mask1 & mask2) =/= 0.U
}
// 判断一条微指令是否被错误预测flush掉
object killedByMispredict{
    def apply(brUpdate: BrUpdateInfo, uop: UOp): Bool = {
        maskMatch(uop.brMask, brUpdate.updateMask.mispredictMask)
    }
    def apply(brUpdate: BrUpdateInfo, uopBrMask: UInt): Bool = {
        maskMatch(uopBrMask, brUpdate.updateMask.mispredictMask)
    }
}

// TODO
// should put common classes here ?
object UpdatedNewUop{
    def apply(uop: UOp, brUpdate: BrUpdateInfo): UOp = {
        val newUop = WireInit(uop)
        newUop.brMask := uop.brMask & (~brUpdate.updateMask.resolveMask).asUInt
        newUop
    }
}
object UpdatedBrMask{
    def apply(uop: UOp, brUpdate: BrUpdateInfo): UInt = {
        uop.brMask & (~brUpdate.updateMask.resolveMask).asUInt
    }
    def apply(uopBrMask: UInt, brUpdate: BrUpdateInfo): UInt = {
        uopBrMask & (~brUpdate.updateMask.resolveMask).asUInt
    }
}
object UpdatedBundle{
    def apply(uop: UOp, brUpdate: BrUpdateInfo): UOp = {
        val newUop = WireInit(uop)
        newUop.brMask := UpdatedBrMask(uop, brUpdate)
        newUop
    }
    def apply[T <: HasUOp](bundle: Valid[T], brUpdate: BrUpdateInfo): Valid[T] = {
        val newBundle = WireInit(bundle)
        newBundle.bits.uop.brMask := UpdatedBrMask(bundle.bits.uop.brMask, brUpdate)
        newBundle.valid := bundle.valid && !killedByMispredict(brUpdate, bundle.bits.uop.brMask)
        newBundle
    }
}

// 小于pc最大的bound倍数
object FloorPCToBound{
    def apply(pc:UInt, bound: Int): UInt = {
        // 双反
        (~((~pc).asUInt | (bound-1).U)).asUInt
    }
}

object genImm{
    def apply(formatImm: UInt, immType: UInt): SInt = {
        val imm = WireInit(0.U(32.W))
        imm := MuxLookup(immType, 0.U)(Seq(
            IMM_U5   -> Cat(0.U(27.W),formatImm(14,10)),
            IMM_U12  -> Cat(0.U(20.W),formatImm(21,10)),
            IMM_S12  -> Cat(Fill(20,formatImm(21)),formatImm(21,10)),
            IMM_S14  -> Cat(Fill(16,formatImm(23)),formatImm(23,10),0.U(2.W)),
            IMM_S16  -> Cat(Fill(14,formatImm(25)),formatImm(25,10),0.U(2.W)),
            IMM_U20  -> Cat(formatImm(24,5),0.U(12.W)),
            IMM_S20  -> Cat(Fill(10,formatImm(24)),formatImm(24,5),0.U(2.W)),
            IMM_S26  -> Cat(Fill(4,formatImm(9)),formatImm(9,0),formatImm(25,10),0.U(2.W)),
            IMM_CSR  -> Cat(0.U(18.W),formatImm(23,10))
        ))
        imm.asSInt
    }
}

object FrontendUtility extends FrontendParam{
    // 双反
    def fetchAlign(pc: UInt): UInt = ~(~pc | (iCacheParam.fetchBytes - 1).U)
    // 取出pc的高位
    def fetchIdx(pc: UInt): UInt = (pc >> log2Ceil(fetchBytes)).asUInt
    // 根据取指包的基地址pc和包内的索引idx还原出具体的指令pc （一条指令4字节，最后两位为0）
    def getPC(pc: UInt, idx:UInt): UInt = Cat(fetchIdx(pc), idx(log2Ceil(fetchWidth)-1, 0), 0.U(2.W))
    def fetchMask(addr: UInt): UInt = {
        // 当前地址是取值pack内的第几条指令
        val idx = addr(log2Ceil(fetchBytes)-1, log2Ceil(instrBytes))
        // 1111 左移 idx 位，得到取值有效掩码（最低位代表pack内的第一条指令有效！）
        (((1<<fetchWidth)-1).U << idx).asUInt(fetchWidth-1, 0)
    }
    // 循环左移
    def rotateLeft(in : UInt, k : Int) = {
        val n = in.getWidth
        Cat(in(n-k-1,0),in(n-1,n-k))
    }
    def nextFetch(addr: UInt): UInt = {
        fetchAlign(addr) + fetchBytes.U
    }

}

object ImplicitCast {
    implicit def uintToBitPat(x: UInt): BitPat = BitPat(x)
}

////////////////// 替换策略
abstract class ReplPolicy(nWays: Int) extends Module {
    require(isPow2(nWays))
    val io = IO(new Bundle {
        val access   = Flipped(Valid(UInt(log2Ceil(nWays).W)))
        val repl_way = Output(UInt(log2Ceil(nWays).W))
    })
}

class PseudoLRU(nWays: Int) extends ReplPolicy(nWays) {
    val tree = RegInit(0.U((nWays - 1).W))

    // update way
    def update_tree(_tree: UInt, access_way: UInt, num_ways: Int): UInt = {
        if (num_ways <= 2) {
            assert(num_ways == 2)
            return !access_way(0)
        } else {
            val left_subtree  = _tree(num_ways - 3, num_ways / 2 - 1)
            val right_subtree = _tree(num_ways / 2 - 2, 0)
            assert(left_subtree.getWidth == num_ways / 2 - 1)
            assert(right_subtree.getWidth == num_ways / 2 - 1)
            val direction = !access_way(log2Ceil(num_ways) - 1)
            return Cat(direction,
                Mux(direction,
                    update_tree(left_subtree, access_way(log2Ceil(num_ways) - 2, 0), num_ways / 2),
                    left_subtree
                ),
                Mux(direction,
                    right_subtree,
                    update_tree(right_subtree, access_way(log2Ceil(num_ways) - 2, 0), num_ways / 2)
                )
            )
        }
    }
    tree := Mux(io.access.valid, update_tree(tree, io.access.bits, nWays), tree)

    // replace way
    def calc_replace_way(_tree: UInt, num_ways: Int): UInt = {
        if (num_ways <= 2) {
            assert(num_ways == 2)
            return _tree(0)
        } else {
            val left_subtree  = _tree(num_ways - 3, num_ways / 2 - 1)
            val right_subtree = _tree(num_ways / 2 - 2, 0)
            assert(left_subtree.getWidth == num_ways / 2 - 1)
            assert(right_subtree.getWidth == num_ways / 2 - 1)
            val prio = _tree(num_ways - 2)
            return Cat(prio, Mux(prio,
                calc_replace_way(right_subtree, num_ways / 2),
                calc_replace_way(left_subtree , num_ways / 2)
            ))
        }
    }
    io.repl_way := calc_replace_way(tree, nWays)
}
